#!/usr/bin/env python
# -*- coding:utf-8 -*-
# @Date : 2024/10/10
# @Author : shiyou pan

from typing import Literal

import numpy as np
import pandas as pd
import requests

from macros.utils import parse_json, intnx

_cookies = {'wzws_sessionid': 'oGV/6pOBOWMxY2UygDEyMC4yMzkuMTY0LjE0MYJmYzVlZTE=',
            'JSESSIONID': 'yct7q3ofvz7LfQe5YyUTlBhfTPvfNJcEJfV9rrqfLl88HDOMMfl7!-1401403675',
            'u': '6', }
_headers = {'Accept': 'application/json, text/javascript, */*; q=0.01',
            'Accept-Language': 'en-US,en;q=0.9',
            'Connection': 'keep-alive',
            'Referer': 'https://data.stats.gov.cn/easyquery.htm?cn=C01',
            'Sec-Fetch-Dest': 'empty',
            'Sec-Fetch-Mode': 'cors',
            'Sec-Fetch-Site': 'same-origin',
            'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/'
                          '120.0.0.0 Safari/537.36',
            'X-Requested-With': 'XMLHttpRequest',
            'sec-ch-ua': '"Not_A Brand";v="8", "Chromium";v="120", "Google Chrome";v="120"',
            'sec-ch-ua-mobile': '?0',
            'sec-ch-ua-platform': '"Windows"', }


def get_macro_tree(ids: str, freq: Literal['yd', 'nd', 'jd'] = 'yd') -> pd.DataFrame:
    """
    获取宏观指标树

    Parameters
    ----------
    ids: str
        指标代码
    freq: Literal['yd', 'nd', 'jd'] = 'yd'，
        数据频率，可选值
        - 'yd': 月度
        - 'nd': 年度
        - 'jd': 季度

    Returns
    -------
    dsn_: pd.DataFrame
        宏观指标树
    """
    data = {'id': ids,
            'dbcode': f'hg{freq}',
            'wdcode': 'zb',
            'm': 'getTree', }
    response = requests.post('https://data.stats.gov.cn/easyquery.htm',
                             cookies=_cookies, headers=_headers, data=data, verify=False)
    response_text = parse_json(response.text)
    dsn_ = pd.DataFrame(response_text)
    return dsn_


def get_macro_data(ids: str, freq: Literal['yd', 'nd', 'jd'] = 'yd') -> tuple[pd.DataFrame, pd.DataFrame]:
    """
    从国家统计官网获取宏观数据

    Parameters
    ----------
    ids: str
        指标代码
    freq: Literal['yd', 'nd', 'jd'] = 'yd'，
        数据频率，可选值
        - 'yd': 月度
        - 'nd': 年度
        - 'jd': 季度

    Returns
    -------
    node: pd.DataFrame
        指标信息
    dsn: pd.DataFrame
        宏观数据集
    """
    dictq = {'A': '03', 'B': '06', 'C': '09', 'D': '12'}
    dictf = {'nd': 'LAST30', 'jd': 'LAST120', 'yd': 'LAST360'}
    sj = dictf.get(freq)  # 时间参数

    params = {'m': 'QueryData',
              'dbcode': f'hg{freq}',
              'rowcode': 'zb',
              'colcode': 'sj',
              'wds': '[]',
              'dfwds': '[{"wdcode":"zb","valuecode":"%s"},'
                       '{"wdcode":"sj","valuecode":"%s"}]' % (ids, sj),
              'k1': '1702882738974',
              'h': '1', }

    response = requests.get('https://data.stats.gov.cn/easyquery.htm',
                            params=params, cookies=_cookies, headers=_headers, verify=False)
    response_text = parse_json(response.text)

    response_text = response_text.get('returndata')
    node1 = response_text.get('wdnodes')
    node = node1[0].get('nodes')
    node = pd.DataFrame(node)

    dict_node = dict(zip(node['code'], node['cname']))
    node2 = response_text.get('datanodes')
    zb_ = [x.get('code').split('_')[0].split('.')[1] for x in node2]
    sj_ = [x.get('code').split('_')[1].split('.')[1] for x in node2]
    data_ = [float(x['data']['data']) for x in node2]
    dsn = pd.DataFrame(np.array([sj_, zb_, data_]).T, columns=['date', 'varname_en', 'value'])
    dsn['varname_zh'] = dsn['varname_en'].map(dict_node)

    # 日期处理
    if freq == 'yd':
        dsn['date'] = pd.to_datetime(dsn['date'].apply(lambda x: '%s/%s/01' % (x[:4], x[-2:])))
    elif freq == 'jd':
        dsn['date'] = pd.to_datetime(dsn['date'].apply(lambda x: '%s/%s/01' % (x[:4], dictq[x[-1]])))
    elif freq == 'nd':
        dsn['date'] = pd.to_datetime(dsn['date'] + '/12/01')
    dsn['date'] = dsn['date'].astype(str).apply(lambda x: intnx(x, 'm', 0, 'e'))

    return node, dsn


if __name__ == '__main__':
    id1 = 'A0201'
    node1_, dsn1 = get_macro_data(id1, freq='jd')
    print(dsn1)
